from torchvision import transforms, datasets
from torch.utils.data import DataLoader


# data augmentation
data_trans = {
    'train': transforms.Compose([
        transforms.RandomRotation(45),  # between -45 degree and +45 degree
        transforms.CenterCrop(224),  # 224*224
        transforms.RandomHorizontalFlip(p=0.5),
        transforms.RandomVerticalFlip(p=0.5),  # randomly flip with probability 0.5
        transforms.ColorJitter(brightness=0.2, contrast=0.1, saturation=0.1, hue=0.1),
        transforms.RandomGrayscale(p=0.025),  # whether to be converted to gray(3 same features)
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])  # mean and variance
    ]),
    'valid': transforms.Compose([
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])
}


# get data called cifar10
def mydata():
    root = 'D:/gitProject/data/'
    ds = {'train': datasets.CIFAR10(root, train=True, transform=transforms.ToTensor(), download=True),
      'valid': datasets.CIFAR10(root, train=False, transform=transforms.ToTensor(), download=True)}
    b_size = 32
    custom_data_loader = {x: DataLoader(dataset=ds[x], batch_size=b_size, shuffle=True) for x in ['train', 'valid']}
    return custom_data_loader
